#!/usr/bin/python

import os
import re
import fnmatch

# Path to unrar binary
UNRAR = '/Users/valankar/rar/unrar'
# Where active torrents are stored by Azureus
ACTIVE_TORRENT_DIR = '/Users/valankar/Library/Application Support/Azureus/active'
# Where rared movies are
RAR_DIR = '/Users/valankar/Movies'
# Where rared movies should be extracted to
EXTRACT_DIR = '/Volumes/Drobo2/Movies'
MOVIE_DIRS = [ EXTRACT_DIR, '/Volumes/Drobo/Movies' ]

def CleanupTitle(title):
  return title.replace('.', ' ').title()

def GetTitle(path):
  found = re.compile(r'(.*).\d{4}.*').findall(path)
  if found:
    candidate = found[0]
  else:
    candidate = path
  found = re.compile(r'(.*?)[\. ](retail|docu|readnfo|repack|unrated|festival|dc|real|limited|proper|r5|dvd)',
                     re.IGNORECASE).findall(candidate)
  if found:
    return CleanupTitle(found[0][0])
  if candidate != path:
    return CleanupTitle(candidate)
  return None
  
def PruneRars(rars):
  """Only keep the first rar."""
  if rars and 'part' in rars[0]:
    return rars[0:1]
  return rars

def FindFirstRarsFromFilelist(filelist):
  """Matches .rar, .001, etc."""
  rars = [x for x in filelist if x.endswith('.rar') or x.endswith('.RAR')]
  if rars:
    return rars
  rars = [x for x in filelist if x.endswith('.001')]
  if rars:
    return rars
  return []

def GetRars(dir):
  filelist = os.listdir(dir)
  rars = PruneRars(FindFirstRarsFromFilelist(filelist))
  subdirs = [x for x in filelist if os.path.isdir(os.path.join(dir, x))]
  for subdir in subdirs:
    filelist = os.listdir(os.path.join(dir, subdir))
    rars.extend(PruneRars([os.path.join(subdir, x) for x in FindFirstRarsFromFilelist(filelist)]))
  return rars

def GetNfos(dir):
  return fnmatch.filter(os.listdir(dir), '*.[Nn][Ff][Oo]')

def GetAvis(dir):
  dirlist = os.listdir(dir)
  found = fnmatch.filter(dirlist, '*.[Aa][Vv][Ii]')
  found += fnmatch.filter(dirlist, '*.[Mm][Pp]4')
  return found

def CopyFile(src, dest):
  if not os.path.exists(dest):
    print 'Copying %s to %s' % (src, dest)
    os.spawnlp(os.P_WAIT, 'cp', 'cp', src, dest)

def ProcessMovie(movie_dir, to_dir, movie_dirs, title, prompt=True):
  """Should chdir to movie directory before calling this."""
  unrar_success = True
  movie_basename = os.path.basename(movie_dir)
  os.chdir(to_dir)

  # Skip movies that already exist.
  for exist_dir in movie_dirs:
    if os.path.exists(os.path.join(exist_dir, title)):
      print 'Path %s already exists, skipping extraction.' % (title,)
      return

  # Unrar from old dir into new
  if prompt:
    user_title = raw_input('New name for %s [%s] (n to skip):\n' % (movie_basename, title))
    if user_title == 'n':
      return
    if user_title:
      title = user_title
  
  # Figure out which rar file to start with
  rars = GetRars(movie_dir)
  avis = GetAvis(movie_dir)
  print 'Rars: %s, Avis: %s' % (rars, avis)
  
  # Make new dir and unrar/copy into it
  os.mkdir(title)
  os.chdir(title)
  if not rars:
    for avi in avis:
      CopyFile(os.path.join(movie_dir, avi), os.path.join(to_dir, title, avi))

  for rar in rars:
    retval = os.spawnlp(os.P_WAIT, UNRAR, UNRAR, 'x', os.path.join(movie_dir, rar))
    unrar_success = unrar_success and not retval

  # Check for rars that exist in extracted dir (subtitles, etc)
  for rar in GetRars('.'):
    if not os.spawnlp(os.P_WAIT, UNRAR, UNRAR, 'x', rar):
      os.spawnlp(os.P_WAIT, 'rm', 'rm', '-f', rar)
      
  # Get back out to where we were
  os.chdir(to_dir)

  # Make sure files exist in destination dir
  if not os.listdir(title):
    print 'Destination dir %s has no files, removing and moving on.' % (title,)
    os.rmdir(title)
    return

  # Check if any unrars failed
  if not unrar_success:
    print 'Some unrars failed. Removing destination dir %s and moving on.' % (title,)
    assert not title.startswith('/')
    os.spawnlp(os.P_WAIT, 'rm', 'rm', '-rf', title)
    return

  # Copy nfo files  
  for nfo in GetNfos(movie_dir):
    CopyFile(os.path.join(movie_dir, nfo), os.path.join(title, nfo))

def UnrarMovies(from_dir, to_dir, movie_dirs):
  os.chdir(from_dir)
  dirlist = [x for x in os.listdir('.') if os.path.isdir(x)]
  dirlist = [x for x in dirlist if x.count('.') > 1 or x.count(' ') > 1]
  for dir in dirlist:
    title = GetTitle(dir)
    if not title:
      print 'Unknown: %s' % (dir,)
      continue
    ProcessMovie(os.path.join(from_dir, dir), to_dir, movie_dirs, title, prompt=False)


if __name__ == '__main__':
  UnrarMovies(RAR_DIR, EXTRACT_DIR, MOVIE_DIRS)
